package bayes

import org.apache.log4j.{Level, Logger}
import org.apache.spark.ml.evaluation.RegressionEvaluator
import org.apache.spark.ml.recommendation.{ALS, ALSModel}
import org.apache.spark.sql.{DataFrame, SparkSession}

/**
  * An example demonstrating ALS.
  * Run with
  * {{{
  * bin/run-example ml.ALSExample
  * }}}
  */
object ALSExample {

  // 定义评分数据封装类
  case class Rating(userId: Int, movieId: Int, rating: Float, timestamp: Long)

  // 定义数据解析封装函数
  def parseRating(str: String): Rating = {
    val fields = str.split("::")
    assert(fields.size == 4)
    Rating(fields(0).toInt, fields(1).toInt, fields(2).toFloat, fields(3).toLong)
  }


  def main(args: Array[String]) {
    Logger.getLogger("org").setLevel(Level.WARN)
    val spark = SparkSession
      .builder
      .appName("ALSExample")
      .master("local[*]")
      .getOrCreate()
    import spark.implicits._

    // 加载评分数据，并解析为Rating对象
    val ratings = spark.read.textFile("G:\\whale\\doit_recommend\\data\\als\\sample_movielens_ratings.txt")
      .map(parseRating)
      .toDF()

    // 划分训练数据与测试数据集
    val Array(training, test) = ratings.randomSplit(Array(0.8, 0.2))

    // 基于训练数据构建ALS推荐模型
    val als = new ALS()
      .setMaxIter(5)
      .setRegParam(0.01)
      .setUserCol("userId")
      .setItemCol("movieId")
      .setRatingCol("rating")

    // 训练模型（求解方程组）
    val model: ALSModel = als.fit(training)

    // 基于RMSE对test数据集的预测结果进行模型评估
    // 设置冷启动策略为drop，对于新注册用户避免产生空推荐矩阵
    model.setColdStartStrategy("drop")
    // 针对测试数据集进行模型预测
    val predictions: DataFrame = model.transform(test)

    predictions.printSchema()
    predictions.show(10,false)



    // 构造一个评估模型，设置评估指标为RMSE
    val evaluator = new RegressionEvaluator()
      .setMetricName("rmse")  // rmse 均方差
      .setLabelCol("rating")
      .setPredictionCol("prediction")

    // 评估预测结果
    val rmse = evaluator.evaluate(predictions)
    println(s"均方误差 = $rmse")

    // 均方误差 = 1.827332995575392
    //sys.exit(1)

    // 为每个用户推荐top10个物品
    val userRecs: DataFrame = model.recommendForAllUsers(10)

    // 为每个物品推荐top10个用户
    val movieRecs = model.recommendForAllItems(10)

    /*
    // 抽取用户数据子集
    val users = ratings.select(als.getUserCol).distinct().limit(3)
    // 为子集用户推荐top10部电影
    val userSubsetRecs = model.recommendForUserSubset(users, 10)

    // 抽取电影子集
    val movies = ratings.select(als.getItemCol).distinct().limit(3)
    // 为子集电影推荐top10用户
    val movieSubSetRecs = model.recommendForItemSubset(movies, 10)*/

    // 打印各项推荐结果
    /*userRecs.printSchema()
    userRecs.show(10,false)


    movieRecs.printSchema()
    movieRecs.show(10,false)


    userSubsetRecs.printSchema()
    userSubsetRecs.show(10,false)

    movieSubSetRecs.printSchema()
    movieSubSetRecs.show(10,false)*/

    // 保存推荐结果
    userRecs.write.json("G:\\whale\\doit_recommend\\data\\rec_result")


    spark.stop()
  }
}

